from numpy import ndarray
import numpy as np
from scipy.fft import dct, dctn, idctn
from moudels.utills import zigzag_order

class DCT:
    # 归一化
    _norm = None #"ortho"
    _dct_dtype = np.float64
    
    @classmethod
    def dct_transform(cls, image:ndarray):
        """对图像进行8×8分块的DCT变换"""

        dct_image = np.zeros((64, 64, 8, 8), dtype=cls._dct_dtype)

        for i in range(64):
            for j in range(64):
                block = image[i*8:(i+1)*8, j*8:(j+1)*8]
                dct_block = dctn(block, norm=cls._norm)
                dct_image[i, j] = dct_block

        return dct_image

    @classmethod
    def idct_transform(cls, dct_image: ndarray):
        """对8×8分块的DCT图像进行IDCT反变换"""
        h, w = dct_image.shape[0] * 8, dct_image.shape[1] * 8
        # 先设为浮点型用于计算
        idct_image = np.zeros((h, w), dtype=cls._dct_dtype)

        for i in range(64):
            for j in range(64):
                dct_block = dct_image[i, j]
                idct_block = idctn(dct_block, norm=cls._norm)  # 使用 norm=None 不
                idct_image[i * 8:(i + 1) * 8, j * 8:(j + 1) * 8] = idct_block

        # 将结果转换为 uint8，确保像素值在 0-255 范围内
        # idct_image = np.clip(idct_image, 0, 255).astype(np.uint8)
        return idct_image
    
    @classmethod
    def extract_low_freq(cls, dct_block:ndarray)->ndarray:
        """从一个8*8dct分块中提取16个低频系数"""
        low_freq = np.zeros((16), dtype=cls._dct_dtype)
  
        zigzag_indices = zigzag_order()
        for k, (x, y) in enumerate(zigzag_indices):
            low_freq[k] = dct_block[x, y]

        return low_freq.reshape(4,4)
    @classmethod
    def reconstruct_dct_block(cls, dct_block:ndarray, low_freq:ndarray)->ndarray:
        """将16个低频系数放回8*8dct分块"""        
        zigzag_indices = zigzag_order()
        low_freqs = low_freq.flatten()
        for k, (x, y) in enumerate(zigzag_indices):
            dct_block[x, y] = low_freqs[k]
        return dct_block

